import numpy as np
from onnxruntime import InferenceSession, SessionOptions
from onnxruntime_extensions import get_library_path, PyOp, onnx_op

# 1. 定义自定义算子实现
@onnx_op(op_type="mydomain::CustomAdd",  # 注意域名需要与PyTorch导出时一致
         inputs=[PyOp.dt_float, PyOp.dt_float, PyOp.dt_float],  # 三个输入: x, y, scale
         outputs=[PyOp.dt_float])  # 一个输出: (x + y) * scale
def custom_add_ort(x, y, scale, **kwargs):
    """
    自定义加法算子实现
    :param x: 输入张量1
    :param y: 输入张量2
    :param scale: 缩放因子
    :param kwargs: 其他可能的属性（本算子不需要）
    :return: (x + y) * scale
    """
    # 确保所有输入都是NumPy数组
    x = np.asarray(x)
    y = np.asarray(y)
    scale = np.asarray(scale)
    
    # 执行计算: (x + y) * scale
    # 注意：scale可能是标量或与x/y相同形状的张量
    return (x + y) * scale

# 2. 创建支持自定义op的InferenceSession
def create_custom_session(onnx_path):
    # 创建会话选项
    sess_options = SessionOptions()
    
    # 注册扩展库 - 这是关键步骤！
    library_path = get_library_path()
    if library_path:
        sess_options.register_custom_ops_library(library_path)
    
    # 创建支持自定义op的会话
    session = InferenceSession(
        onnx_path,
        sess_options,
        providers=['CPUExecutionProvider']
    )
    return session

# 3. 加载并运行模型
if __name__ == "__main__":
    # 模型路径
    onnx_path = "custom_add_model.onnx"
    
    # 创建推理会话
    session = create_custom_session(onnx_path)
    
    # 准备输入数据（与PyTorch导出时一致）
    input_x = np.array([1.0, 2.0, 3.0], dtype=np.float32)
    input_y = np.array([4.0, 5.0, 6.0], dtype=np.float32)
    
    # 运行推理
    outputs = session.run(
        None,  # 获取所有输出
        {
            "input_x": input_x,
            "input_y": input_y
        }
    )
    result = outputs[0]
    
    print("ONNX Runtime Output:", result)
    print("验证结果:", np.allclose(result, [15., 21., 27.]))
